import re
from datetime import datetime, timedelta

import numpy as np

from components.config import NET_ROBOT_MYSQL_CONFIG
from utils.db.mysqldb import MysqlDB
from utils.db.redisdb import RedisDB
import abc
import concurrent.futures
import warnings

warnings.filterwarnings("ignore")
import pandas as pd

from loguru import logger

from utils.es_company_tools import EsCompanyTools


class BaseTaxIllegalClean(object):
    def __init__(self, origin_table, target_table, redis_task_keys, origin_table_fields=None):
        if origin_table_fields is None:
            origin_table_fields = []
        self.temporary_data = {}
        self.origin_table = origin_table  # 来源库表
        self.target_table = target_table  # 目标库表
        self.origin_table_fields = origin_table_fields  # 来源库表字段
        self.redis_task_keys = redis_task_keys  # redis 任务表名
        self.mysql_db_origin = MysqlDB(
            ip=NET_ROBOT_MYSQL_CONFIG["MYSQL_IP"],
            port=NET_ROBOT_MYSQL_CONFIG["MYSQL_PORT"],
            db=NET_ROBOT_MYSQL_CONFIG["MYSQL_DB"],
            user_name=NET_ROBOT_MYSQL_CONFIG["MYSQL_USER_NAME"],
            user_pass=NET_ROBOT_MYSQL_CONFIG["MYSQL_USER_PASS"],
        )
        self.mysql_db_target = self.mysql_db_origin
        self.redis_db = RedisDB()
        self.redis_task_keys_error = self.redis_task_keys + ":clean:error"
        logger.info(
            f"\n {self.redis_db} redis_task_keys: {self.redis_task_keys}\nredis_task_keys_error:" f" {self.redis_task_keys_error}"
        )

    def get_data(self, id_list: list) -> pd.DataFrame:
        """
        获取原始数据
        """
        data = []
        id_list_str = ",".join(f"'{i}'" for i in id_list)
        if id_list_str:
            _table_fields = ",".join(self.origin_table_fields) if self.origin_table_fields else "*"
            sql = "select {} from {} where id in ({})".format(_table_fields, self.origin_table, id_list_str)
            data = self.mysql_db_origin.find(sql, to_json=True)
        return pd.DataFrame(data)

    def time_wash(self, start_time, end_time, batch_size=1000, worker=1):
        """
        定时清洗
        """
        self._push_task(start_time=start_time, end_time=end_time, batch_size=batch_size)  # 推送任务
        with concurrent.futures.ThreadPoolExecutor(max_workers=worker) as executor:
            executor.submit(self._run_task)  # 执行任务

    @abc.abstractmethod
    def run(self, id_list):
        """
        实际清洗方法 需要重写
        """
        df_src = self.get_data(id_list)  # 原始数据
        df_src = self.process_field_test(df_src)  # 一层数据流转处理
        if not df_src.empty:
            # 保存
            df_src = df_src.drop(columns=["id"])
            df_src.replace({np.NAN: None}, inplace=True)
            datas = df_src.to_dict("records")
            update_columns = []  # 要更新的字段 按唯一索引更新
            self.mysql_db_target.add_batch_smart(self.target_table, datas, update_columns=update_columns)

    def process_field_test(self, df_pro: pd.DataFrame) -> pd.DataFrame:
        """
        具体的清洗逻辑
        """
        # 删除创建时间和更新时间
        df_pro = df_pro.drop(columns=["create_time", "update_time"])

        # 插入两列数据
        df_pro.loc[:, "data_sources"] = "国家税务总局"
        df_pro.loc[:, "publish_date"] = 0
        for index, row in df_pro.iterrows():

            if row["taxpayer_id"] == "":
                company_name = row["company_name"]
                df_pro["taxpayer_id"][index] = (
                    EsCompanyTools()
                    .get_company_name_or_taxpayer_id([company_name])["taxpayer_id_mapping"]
                    .get(company_name)
                )

            df_pro["publish_date"][index] = self.publish_date(
                row["year"], row["illegal_facts"]["content"] if row["illegal_facts"] else ""
            )
            df_pro["year"][index] = str(row["year"])[:4]

        return df_pro

    def _push_task(self, start_time: str, end_time: str, batch_size=1000) -> None:
        """
        推送任务
        """
        sql = "select id from {} where update_time >= '{}' and update_time < '{}' ".format(
            self.origin_table, start_time, end_time
        )
        logger.info(f"{self.origin_table} 开始查询并推送任务，sql : {sql}")
        id_list = self.mysql_db_origin.find(sql, to_json=True)
        id_list = [i["id"] for i in id_list]
        for i in range(0, len(id_list), batch_size):
            sub_ids = id_list[i:i + batch_size]
            self.redis_db.sadd(self.redis_task_keys, str(sub_ids))  # 推送任务

    def _run_task(self):
        """
        消费任务 错误的任务记录 可开多线程
        :return:
        """
        max_retry = 3
        while True:
            logger.info(f"{self.origin_table} 开始消费任务，剩余 {self.redis_db.sget_count(self.redis_task_keys)}")
            data_str = self.redis_db.sget(self.redis_task_keys, is_pop=True)
            if data_str:
                id_list = eval(data_str[0])  # id 列表
                logger.debug(f"当前任务: {str(id_list)[:20]} ...")
                for _ in range(1, max_retry + 1):
                    try:
                        self.run(id_list)  # 清洗入库
                        break
                    except KeyboardInterrupt as e:
                        self.redis_db.sadd(self.redis_task_keys, str(id_list))
                    except Exception as e:
                        logger.error(f"{self.origin_table} 任务出错，错误信息(重新推送)：{e} {str(id_list)[:50]}")
                else:
                    # 重试超过次数 不再重试
                    logger.error(f"{self.origin_table} 任务出错，超过重试次数) {str(id_list)[:50]}")
                    self.redis_db.sadd(self.redis_task_keys_error, str(id_list))
            else:
                logger.info("我的任务完成了")
                break

    def publish_date(self, year, illegal_facts):
        if len(str(year)) == 8:
            if datetime(int(str(year)[:4]), int(str(year)[4:6]), int(str(year)[6:8])) > datetime.now():
                return None
            return datetime(int(str(year)[:4]), int(str(year)[4:6]), int(str(year)[6:8]))
        elif len(str(year)) == 4:
            if datetime(year, 1, 1) > datetime.now():
                return None
            return datetime(year, 1, 1)
        elif year == 0:
            date_list = re.findall(r"(\d{4})年", str(illegal_facts))
            if len(date_list) == 0:
                return None
            elif len(date_list) >= 2:
                if datetime(int(date_list[1]), 1, 1) > datetime.now():
                    return None
                return datetime(int(date_list[1]), 1, 1)
            else:
                return None


def run_clean():
    test = BaseTaxIllegalClean(
        origin_table="net_tax_illegal",
        target_table="clean_tax_illegal",
        redis_task_keys="黑名单重大税务案件清洗:wash:task",
    )
    test.time_wash(
        start_time=str(datetime.now().date() - timedelta(days=5)),
        end_time=str(datetime.now().date() + timedelta(days=1)),
        batch_size=100,
        worker=1,
    )


if __name__ == "__main__":
    run_clean()
